package me.zhengjie.stock.pipeline;

import com.alibaba.fastjson.JSONArray;
import com.alibaba.fastjson.JSONObject;
import com.baomidou.mybatisplus.core.conditions.query.LambdaQueryWrapper;
import lombok.AllArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import me.zhengjie.stock.domain.TblStockInfo;
import me.zhengjie.stock.service.TblStockInfoService;
import org.springframework.stereotype.Component;
import us.codecraft.webmagic.ResultItems;
import us.codecraft.webmagic.Task;
import us.codecraft.webmagic.pipeline.Pipeline;

import java.util.ArrayList;
import java.util.Date;
import java.util.List;

/**
 * @author Aaron
 * @date 2020/12/19
 */
@Slf4j
@Component
@AllArgsConstructor
public class StockInfoPipeline implements Pipeline {

    private final TblStockInfoService tblStockInfoService;

    @Override
    public void process(ResultItems resultItems, Task task) {
        if (resultItems.isSkip()) {
            return;
        }
        JSONArray dataList = resultItems.get("dataList");
        if (dataList != null && !dataList.isEmpty()) {
            int size = dataList.size();
            List<TblStockInfo> tblStockInfoList = new ArrayList<>();
            for (int i = 0; i < size; i++) {
                JSONObject stockInfoJson = dataList.getJSONObject(i);
                String code = stockInfoJson.getString("f12");
                String name = stockInfoJson.getString("f14");
                TblStockInfo tblStockInfo = new TblStockInfo();
                tblStockInfo.setName(name);
                tblStockInfo.setCode(code);
                tblStockInfo.setIndustry(stockInfoJson.getString("f100"));
                tblStockInfo.setCity(stockInfoJson.getString("f102"));
                tblStockInfo.setConcept(stockInfoJson.getString("f103"));
                tblStockInfo.setMarket(getMarket(code));
                tblStockInfo.setOpenStatus(1);
                tblStockInfo.setGmtModified(new Date());
                if (name.contains("PT") || name.contains("ST") || name.endsWith("退")) {
                    //ST、退市等股票设置为0
                    tblStockInfo.setIndustry("ST/退");
                    tblStockInfo.setOpenStatus(0);
                }
                if (code.startsWith("30")) {
                    //创业板股票设置为0
                    tblStockInfo.setIndustry("创业板");
                    tblStockInfo.setOpenStatus(0);
                }
                if (code.startsWith("688")) {
                    //科创板股票设置为0
                    tblStockInfo.setIndustry("科创板");
                    tblStockInfo.setOpenStatus(0);
                }
                if (code.startsWith("43") || code.startsWith("83")) {
                    //科创板股票设置为0
                    tblStockInfo.setIndustry("北交所");
                    tblStockInfo.setOpenStatus(0);
                }
                LambdaQueryWrapper<TblStockInfo> queryWrapper = new LambdaQueryWrapper<>();
                queryWrapper.eq(TblStockInfo::getCode, code);
                queryWrapper.eq(TblStockInfo::getName, name);
                TblStockInfo stockInfoDb = tblStockInfoService.getOne(queryWrapper);
                if (stockInfoDb != null) {
                    tblStockInfo.setId(stockInfoDb.getId());
                    tblStockInfoService.updateById(tblStockInfo);
                    log.info("save StockInfoPipeline to db fail! code:{}, name:{} exited", code, name);
                    continue;
                }
                tblStockInfo.setGmtCreate(new Date());
                tblStockInfoList.add(tblStockInfo);
                //保存数据
                log.info("save StockInfoPipeline to db success! code:{}, name:{}", code, name);
            }
            tblStockInfoService.saveBatch(tblStockInfoList);
            log.info("save StockInfoPipeline to Mysql success! size:{}", size);
        }
    }

    private Integer getMarket(String code) {
        //市场，0 深证，1 上证
        return code.startsWith("6") ? 1 : 0;
    }

}
